// Virtual memory management implementation
// Author: $USER_NAME$
// Date: $DATE$
// Description: 实现虚拟内存管理功能，包括页表操作、地址转换、内存分配与释放、
// 进程地址空间管理等核心虚拟内存功能

#include "param.h"
#include "types.h"
#include "memlayout.h"
#include "elf.h"
#include "riscv.h"
#include "defs.h"
#include "fs.h"
#include "spinlock.h"
#include "proc.h"

/*
 * the kernel's page table.
 */
pagetable_t kernel_pagetable;

extern char etext[];  // kernel.ld sets this to end of kernel code.

extern char trampoline[]; // trampoline.S

void print(pagetable_t);

/*
 * create a direct-map page table for the kernel and
 * turn on paging. called early, in supervisor mode.
 * the page allocator is already initialized.
 */
// 为内核创建直接映射页表并启用分页
// 在早期以监管模式调用，页面分配器已初始化
void
kvminit()
{
  kernel_pagetable = (pagetable_t) kalloc();
  memset(kernel_pagetable, 0, PGSIZE);

  // uart registers
  kvmmap(UART0, UART0, PGSIZE, PTE_R | PTE_W);

  // virtio mmio disk interface 0
  kvmmap(VIRTION(0), VIRTION(0), PGSIZE, PTE_R | PTE_W);

  // virtio mmio disk interface 1
  kvmmap(VIRTION(1), VIRTION(1), PGSIZE, PTE_R | PTE_W);

  // CLINT
  kvmmap(CLINT, CLINT, 0x10000, PTE_R | PTE_W);

  // PLIC
  kvmmap(PLIC, PLIC, 0x400000, PTE_R | PTE_W);

  // map kernel text executable and read-only.
  kvmmap(KERNBASE, KERNBASE, (uint64)etext-KERNBASE, PTE_R | PTE_X);

  // map kernel data and the physical RAM we'll make use of.
  kvmmap((uint64)etext, (uint64)etext, PHYSTOP-(uint64)etext, PTE_R | PTE_W);

  // map the trampoline for trap entry/exit to
  // the highest virtual address in the kernel.
  kvmmap(TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_R | PTE_X);
}

// 将硬件页表寄存器切换到内核页表并启用分页
void
kvminithart()
{
  w_satp(MAKE_SATP(kernel_pagetable));
  sfence_vma();
}

// 返回页表pagetable中对应虚拟地址va的PTE地址
// 如果alloc!=0，则创建所需的页表页
//
// risc-v Sv39方案有三级页表页，页表页包含512个64位PTE
// 64位虚拟地址分为以下字段：
//   39..63 -- 必须为零
//   30..38 -- 9位二级索引
//   21..39 -- 9位一级索引
//   12..20 -- 9位零级索引
//    0..12 -- 页内12位字节偏移
static pte_t *
walk(pagetable_t pagetable, uint64 va, int alloc)
{
  if(va >= MAXVA)
    panic("walk");

  for(int level = 2; level > 0; level--) {
    pte_t *pte = &pagetable[PX(level, va)];
    if(*pte & PTE_V) {
      pagetable = (pagetable_t)PTE2PA(*pte);
    } else {
      if(!alloc || (pagetable = (pde_t*)kalloc()) == 0)
        return 0;
      memset(pagetable, 0, PGSIZE);
      *pte = PA2PTE(pagetable) | PTE_V;
    }
  }
  return &pagetable[PX(0, va)];
}

// 查找虚拟地址，返回物理地址，未映射则返回0
// 只能用于查找用户页面
uint64
walkaddr(pagetable_t pagetable, uint64 va)
{
  pte_t *pte;
  uint64 pa;

  if(va >= MAXVA)
    return 0;

  pte = walk(pagetable, va, 0);
  if(pte == 0)
    goto lzac;
  if((*pte & PTE_V) == 0)
    goto lzac;
  if((*pte & PTE_U) == 0)
    goto lzac;
  pa = PTE2PA(*pte);

  if(0) {
  lzac:
    if ((pa = lazyalloc(myproc(), va)) <= 0)
      pa = 0;                                                                                                                                   
  }                                                                                                                                                      
      
  return pa;
}

// 为内核页表添加映射
// 仅在引导时使用，不刷新TLB或启用分页
void
kvmmap(uint64 va, uint64 pa, uint64 sz, int perm)
{
  if(mappages(kernel_pagetable, va, sz, pa, perm) != 0)
    panic("kvmmap");
}

// 将内核虚拟地址转换为物理地址
// 仅需要用于栈上的地址，假设va是页面对齐的
uint64
kvmpa(uint64 va)
{
  uint64 off = va % PGSIZE;
  pte_t *pte;
  uint64 pa;
  
  pte = walk(kernel_pagetable, va, 0);
  if(pte == 0)
    panic("kvmpa");
  if((*pte & PTE_V) == 0)
    panic("kvmpa");
  pa = PTE2PA(*pte);
  return pa+off;
}

// 为从va开始的虚拟地址创建PTE，这些地址引用从pa开始的物理地址
// va和size可能不是页面对齐的
// 返回值: 成功返回0，walk()无法分配所需页表页返回-1
int
mappages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm)
{
  uint64 a, last;
  pte_t *pte;

  a = PGROUNDDOWN(va);
  last = PGROUNDDOWN(va + size - 1);
  for(;;){
    if((pte = walk(pagetable, a, 1)) == 0)
      return -1;
    if(*pte & PTE_V)
      ;
//  panic("remap");
    *pte = PA2PTE(pa) | perm | PTE_V;
    if(a == last)
      break;
    a += PGSIZE;
    pa += PGSIZE;
  }
  return 0;
}

// 从页表中移除映射，给定范围内的映射必须存在
// 可选择是否释放物理内存
void
uvmunmap(pagetable_t pagetable, uint64 va, uint64 size, int do_free)
{
  uint64 a, last;
  pte_t *pte;
  uint64 pa;

  a = PGROUNDDOWN(va);
  last = PGROUNDDOWN(va + size - 1);
  if (a > last)
    return;
  for(;;){
    if ((pte = walk(pagetable, a, 0)) == 0) {
      goto out;
    }

    if ((*pte & PTE_V) == 0) {
      goto out;
    }
    if(PTE_FLAGS(*pte) == PTE_V)
      panic("uvmunmap: not a leaf");
    if(do_free){
      pa = PTE2PA(*pte);
      kfree((void*)pa);
    }

    *pte = 0;
  out:    
    if(a == last)
      break;
    a += PGSIZE;
    pa += PGSIZE;
  }

}

// 创建空的用户页表
// 返回值: 成功返回页表指针，失败则panic
pagetable_t
uvmcreate()
{
  pagetable_t pagetable;
  pagetable = (pagetable_t) kalloc();
  if(pagetable == 0)
    panic("uvmcreate: out of memory");
  memset(pagetable, 0, PGSIZE);
  return pagetable;
}

// 将用户initcode加载到页表地址0处，用于第一个进程
// sz必须小于一页
void
uvminit(pagetable_t pagetable, uchar *src, uint sz)
{
  char *mem;

  if(sz >= PGSIZE)
    panic("inituvm: more than a page");
  mem = kalloc();
  memset(mem, 0, PGSIZE);
  mappages(pagetable, 0, PGSIZE, (uint64)mem, PTE_W|PTE_R|PTE_X|PTE_U);
  memmove(mem, src, sz);
}

// 分配PTE和物理内存，将进程从oldsz增长到newsz
// newsz不需要页面对齐
// 返回值: 新的大小，错误返回0
uint64
uvmalloc(pagetable_t pagetable, uint64 oldsz, uint64 newsz)
{
  char *mem;
  uint64 a;

  if(newsz < oldsz)
    return oldsz;

  oldsz = PGROUNDUP(oldsz);
  a = oldsz;
  for(; a < newsz; a += PGSIZE){
    mem = kalloc();
    if(mem == 0){
      uvmdealloc(pagetable, a, oldsz);
      return 0;
    }
    memset(mem, 0, PGSIZE);
    if(mappages(pagetable, a, PGSIZE, (uint64)mem, PTE_W|PTE_X|PTE_R|PTE_U) != 0){
      kfree(mem);
      uvmdealloc(pagetable, a, oldsz);
      return 0;
    }
  }
  return newsz;
}

// 释放用户页面，将进程大小从oldsz减少到newsz
// oldsz和newsz不需要页面对齐，newsz也不需要小于oldsz
// oldsz可以大于实际进程大小
// 返回值: 新的进程大小
uint64
uvmdealloc(pagetable_t pagetable, uint64 oldsz, uint64 newsz)
{
  if(newsz >= oldsz)
    return oldsz;

  uint64 newup = PGROUNDUP(newsz);
  if(newup < PGROUNDUP(oldsz))
    uvmunmap(pagetable, newup, oldsz - newup, 1);

  return newsz;
}

// 递归释放页表页面，所有叶映射必须已经被移除
static void
freewalk(pagetable_t pagetable)
{
  // there are 2^9 = 512 PTEs in a page table.
  for(int i = 0; i < 512; i++){
    pte_t pte = pagetable[i];
    if((pte & PTE_V) && (pte & (PTE_R|PTE_W|PTE_X)) == 0){
      // this PTE points to a lower-level page table.
      uint64 child = PTE2PA(pte);
      freewalk((pagetable_t)child);
      pagetable[i] = 0;
    } else if(pte & PTE_V){
      panic("freewalk: leaf");
    }
  }
  kfree((void*)pagetable);
}

// 释放用户内存页面，然后释放页表页面
void
uvmfree(pagetable_t pagetable, uint64 sz)
{
  uvmunmap(pagetable, 0, sz, 1);
  freewalk(pagetable);
}

// 给定父进程的页表，将其内存复制到子进程的页表中
// 复制页表和物理内存
// 返回值: 成功返回0，失败返回-1，失败时释放已分配的页面
int
uvmcopy(pagetable_t old, pagetable_t new, uint64 sz)
{
  pte_t *pte;
  uint64 pa, i;
  uint flags;
  char *mem;

  for(i = 0; i < sz; i += PGSIZE){
    if ((pte = walk(old, i, 0)) == 0) {
      continue;
    }
    if ((*pte & PTE_V) == 0) {
      continue;
    }
    pa = PTE2PA(*pte);
    flags = PTE_FLAGS(*pte);
    if((mem = kalloc()) == 0)
      goto err;
    memmove(mem, (char*)pa, PGSIZE);
    if(mappages(new, i, PGSIZE, (uint64)mem, flags) != 0){
      kfree(mem);
      goto err;
    }
  }
  return 0;

 err:
  uvmunmap(new, 0, i, 1);
  return -1;
}

// 标记PTE对用户访问无效，exec用于用户栈保护页面
void
uvmclear(pagetable_t pagetable, uint64 va)
{
  pte_t *pte;
  
  pte = walk(pagetable, va, 0);
  if(pte == 0)
    panic("uvmclear");
  *pte &= ~PTE_U;
}

// 从内核复制到用户空间
// 将len字节从src复制到给定页表中的虚拟地址dstva
// 返回值: 成功返回0，错误返回-1
int
copyout(pagetable_t pagetable, uint64 dstva, char *src, uint64 len)
{
  uint64 n, va0, pa0;

  while(len > 0){
    va0 = PGROUNDDOWN(dstva);
    pa0 = walkaddr(pagetable, va0);
    if(pa0 == 0) {
      if (dstva >= myproc()->sz) {
        return -1;
      }
      char *mem = kalloc();
      pa0= (uint64)mem;
      memset(mem, 0, PGSIZE);
      mappages(pagetable, va0, PGSIZE, pa0, PTE_W|PTE_X|PTE_R|PTE_U);
    }
    n = PGSIZE - (dstva - va0);
    if(n > len)
      n = len;
    memmove((void *)(pa0 + (dstva - va0)), src, n);

    len -= n;
    src += n;
    dstva = va0 + PGSIZE;
  }
  return 0;
}

// 从用户空间复制到内核
// 将len字节复制到dst，从给定页表中的虚拟地址srcva
// 返回值: 成功返回0，错误返回-1
int
copyin(pagetable_t pagetable, char *dst, uint64 srcva, uint64 len)
{
  uint64 n, va0, pa0;

  while(len > 0){
    va0 = PGROUNDDOWN(srcva);
    pa0 = walkaddr(pagetable, va0);
    if(pa0 == 0) {
      if (srcva >= myproc()->sz) {
        return -1;
      }
      char *mem = kalloc();
      pa0 = (uint64)mem;
      memset(mem, 0, PGSIZE);
      mappages(pagetable, va0, PGSIZE, pa0, PTE_W|PTE_X|PTE_R|PTE_U);
    }
    n = PGSIZE - (srcva - va0);
    if(n > len)
      n = len;
    memmove(dst, (void *)(pa0 + (srcva - va0)), n);

    len -= n;
    dst += n;
    srcva = va0 + PGSIZE;
  }
  return 0;
}

// 从用户空间复制以null结尾的字符串到内核
// 将字节复制到dst，从给定页表中的虚拟地址srcva，直到'\0'或max
// 返回值: 成功返回0，错误返回-1
int
copyinstr(pagetable_t pagetable, char *dst, uint64 srcva, uint64 max)
{
  uint64 n, va0, pa0;
  int got_null = 0;

  while(got_null == 0 && max > 0){
    va0 = PGROUNDDOWN(srcva);
    pa0 = walkaddr(pagetable, va0);
    if(pa0 == 0)
      return -1;
    n = PGSIZE - (srcva - va0);
    if(n > max)
      n = max;

    char *p = (char *) (pa0 + (srcva - va0));
    while(n > 0){
      if(*p == '\0'){
        *dst = '\0';
        got_null = 1;
        break;
      } else {
        *dst = *p;
      }
      --n;
      --max;
      p++;
      dst++;
    }

    srcva = va0 + PGSIZE;
  }
  if(got_null){
    return 0;
  } else {
    return -1;
  }
}

// 懒分配物理内存
uint64 lazyalloc(struct proc * p, uint64 va){
#define lazyalloc_debug 0
#define lazyalloc_warn(info) { \
    printf("lazyalloc(): %s", info); \
    printf("             scause %p pid=%d\n", r_scause(), p->pid); \
    printf("             sepc=%p stval=%p\n", r_sepc(), r_stval()); \
  }                                                                                                                                                        
  if(va >= p->sz || va < PGROUNDUP(p->tf->sp)){                                                                                              
#if lazyalloc_debug                                                                                                                             
    lazyalloc_warn("vm addr higher then any allocated with sbrk\n");                                                                              
#endif                                                                                                                                             
    return 0;                                                                                                                                       
  }                                                                                                                                                 
  char * mem;                                                                                                                                       
  uint64 a = PGROUNDDOWN(va);                                                                                                                       
  mem = kalloc();                                                                                                                                   
  if(mem == 0){                                                                                                                                     
#if lazyalloc_debug                                                                                                                             
    lazyalloc_warn("kalloc() == 0\n");                                                                                                            
#endif                                                                                                                                             
    return 0;                                                                                                                                       
  }                                                                                                                                                 
  memset(mem, 0, PGSIZE);                                                                                                                           
  if(mappages(p->pagetable, a, PGSIZE, (uint64)mem, PTE_W|PTE_X|PTE_R|PTE_U) != 0){                                                                    
#if lazyalloc_debug                                                                                                                             
    lazyalloc_warn("mappages() != 0\n");                                                                                                          
#endif                                                                                                                                             
    kfree(mem);                                                                                                                                     
    return 0;                                                                                                                                       
  }                                                                                                                                                    
  
  return (uint64)mem;                                                                                                                               
}